# -*- coding: utf-8 -*-
"""
=============================================================
HHL (cqlib version) — Qiskit-like class & result container
=============================================================

依赖：
- circuits.algorithm_circuit.Algorithm_Circuit as Circuit
- cqlib.simulator.StatevectorSimulator
- algorithm.quantum_phase_estimator:
    hhl_pe_encode, hhl_pe_uncompute, pick_t_from_lambda_max, rough_lambda_max_via_eigs
- algorithm.Hamiltonian_simulator:
    build_pauli_evolution_oracle_from_matrix
- algorithm.state_preparation.prepare_state_from_vector  (若不存在则 2x2 降级)

功能：
- HHL.construct_circuit() 构造电路
- HHL.run() 在 statevector 模式下模拟并产出 HHLResult
- 从量子解自动恢复“经典尺度”的解向量
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List, Tuple
import numpy as np
import math

# cqlib
from circuits.algorithm_circuit import Algorithm_Circuit as Circuit
from cqlib.simulator import StatevectorSimulator

from algorithm.quantum_phase_estimator import (
    hhl_pe_encode, hhl_pe_uncompute,
    pick_t_from_lambda_max, rough_lambda_max_via_eigs,
)
from algorithm.Hamiltonian_simulator import (
    build_pauli_evolution_oracle_from_matrix,
)

from algorithm.lookup_rotation import LookupRotation

from algorithm.quantum_state_prep import prepare_state_from_vector
# ---------------------------
# 结果对象
# ---------------------------
@dataclass
class HHLResult:
    solution: Optional[np.ndarray] = None         # 经典尺度解
    output: Optional[np.ndarray] = None           # 量子方向解（归一）
    probability_result: Optional[float] = None    # anc=1 成功概率
    matrix: Optional[np.ndarray] = None
    vector: Optional[np.ndarray] = None
    circuit_info: Dict[str, Any] = field(default_factory=dict)


# ---------------------------
# 工具：统一态矢量为一维复 ndarray
# ---------------------------
def _coerce_statevector_to_array(sv_like, total_n: int) -> np.ndarray:
    """
    将 statevector 统一成长度 2**total_n 的复数组。
    支持：
      - dict：键为二进制字符串（'0101'）或整数索引
    """
    N = 1 << int(total_n)

    # dict: {bitstring or int: amplitude}
    arr = np.zeros(N, dtype=complex)
    for k, v in sv_like.items():
        i = int(k, 2) if isinstance(k, str) else int(k)
        arr[i] = complex(v)
    return arr


# ---------------------------
# HHL 主类
# ---------------------------
class HHL:
    def __init__(
        self,
        matrix: np.ndarray,
        vector: np.ndarray,
        m_phase: int = 4,
        trotter_slices: int = 2,
        truncate_powerdim: bool = False,
        truncate_hermitian: bool = False,
        orig_size: Optional[int] = None,
    ):
        """
        Args:
            matrix, vector     : 线性系统 Ax=b
            m_phase            : 相位寄存器比特数
            trotter_slices     : 演化分片
            truncate_powerdim  : 若输入非 2^n 维，是否扩到 2^n，完后再截回
            truncate_hermitian : 若输入非厄米，是否扩到厄米块矩阵，完后再截回
            orig_size          : 截回原尺寸需要的原始维度
        """
        A = np.array(matrix, dtype=float)
        b = np.array(vector, dtype=float).reshape(-1)
        if A.shape[0] != A.shape[1]:
            raise ValueError("matrix must be square")
        if A.shape[0] != b.size:
            raise ValueError("vector dimension mismatch")

        self._orig_A = A.copy()
        self._orig_b = b.copy()
        self._truncate_powerdim = truncate_powerdim
        self._truncate_hermitian = truncate_hermitian
        self._orig_size = orig_size or b.size

        # 处理到 2^n & 厄米（若需要）
        A2, b2 = A, b
        if truncate_powerdim and (math.log2(A.shape[0]) % 1 != 0):
            A2, b2 = self.expand_to_powerdim(A2, b2)
        if truncate_hermitian and not np.allclose(A2, A2.conj().T):
            A2, b2 = self.expand_to_hermitian(A2, b2)
        # 基本校验
        n_qubits_work = int(round(math.log2(A2.shape[0])))
        if 2 ** n_qubits_work != A2.shape[0]:
            raise ValueError("matrix dimension must be 2^n after resize")
        if not np.allclose(A2, A2.conj().T):
            raise ValueError("matrix must be Hermitian after resize")

        self.A = A2
        self.b = b2 / np.linalg.norm(b2)  # 量子态制备需要
        self.n_work = n_qubits_work
        self.m_phase = int(m_phase)
        self.trotter_slices = int(trotter_slices)

        # 内部缓存
        self._qc: Optional[Circuit] = None
        self._reg: Optional[Dict[str, List[int]]] = None
        self._t: Optional[float] = None

    # ---------- 扩维工具 ----------
    @staticmethod
    def expand_to_powerdim(matrix: np.ndarray, vector: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        dim = matrix.shape[0]
        next_n = int(np.ceil(np.log2(dim)))
        newA = np.eye(2 ** next_n, dtype=float)
        newA[:dim, :dim] = matrix
        newb = np.zeros(2 ** next_n, dtype=float)
        newb[:vector.size] = vector
        return newA, newb

    @staticmethod
    def expand_to_hermitian(matrix: np.ndarray, vector: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        d = matrix.shape[0]
        A = np.zeros((2 * d, 2 * d), dtype=float)
        A[:d, d:] = matrix.conj().T
        A[d:, :d] = matrix
        b = np.zeros(2 * d, dtype=float)
        b[:d] = vector.conj()
        b[d:] = vector
        return A, b

    # ---------- 构造电路 ----------
    def construct_circuit(self, measurement: bool = False) -> Circuit:
        """
        构造完整 HHL 电路：
          - |b> 制备到 work 寄存器
          - QPE 编码到 phase 寄存器
          - 倒数旋转（使用 LookupRotation）
          - QPE 回滚
        """
        A = self.A
        b = self.b
        n_work = self.n_work
        m_phase = self.m_phase

        # 受控演化 & 时间步
        lam_max = rough_lambda_max_via_eigs(A)
        t = pick_t_from_lambda_max(lam_max, target_phase=0.35, safety=0.95)
        evo = build_pauli_evolution_oracle_from_matrix(
            A, method='suzuki2', num_slices=self.trotter_slices
        )

        # ============ 关键：为 LookupRotation 预留位 ============
        # LookupRotation 会在 inreg 之后分配：
        #   work, msq, anc  共 3 个
        #   以及 aux_needed = (len(inreg) - 1) 个辅助位
        aux_needed_lookup = max(0, m_phase - 1)
        extra_lookup = 3 + aux_needed_lookup

        total_qubits = m_phase + n_work + extra_lookup
        qc = Circuit(total_qubits)

        # 索引分配
        work = list(range(0, n_work))                       
        phase = list(range(n_work, n_work + m_phase))                

        # |b> 制备到 work 寄存器
        prep, _, _, _ = prepare_state_from_vector(list(b))
        # 注意：compose 的第二个参数为“目标量子位索引列表”，把制备电路嵌到 work 上
        qc = qc.compose(prep, work)

        # QPE 编码
        hhl_pe_encode(qc, phase_qubits=phase, work_qubits=work, evolution_oracle=evo, t=t)

        # 倒数旋转（LookupRotation）
        lam_vals = np.linalg.eigvalsh(A)
        lam_min = float(np.min(np.abs(lam_vals)))

        # 创建并附加查表倒数模块；它将使用现有 qc，并在 phase 后续索引处添加 work/msq/anc/aux
        self._lookup = LookupRotation(
            pat_length=None,
            subpat_length=None,
            negative_evals=False,
            evo_time=t,
            lambda_min=lam_min,
        )
        qc = self._lookup.construct_circuit(mode="circuit", inreg=phase, circuit=qc)

        # QPE 回滚
        hhl_pe_uncompute(qc, phase_qubits=phase, work_qubits=work, evolution_oracle=evo, t=t)

        if measurement:
            # 此处不加测量门；只在结果提取时用 anc=1 的幅度
            pass

        # 从 lookup 拿到 anc 索引（用于后续从 statevector 提取）
        anc = int(self._lookup._anc)  # type: ignore[attr-defined]
        # 记录（也把 lookup 的辅助位池记下来，运行时计算总位数用得到）
        aux_lookup = list(getattr(self._lookup, "_aux_pool", []))  # type: ignore[attr-defined]

        self._qc = qc
        self._reg = {"work": work, "phase": phase, "anc": [anc], "aux": aux_lookup}
        return qc

    # ---------- 从 statevector 提取解 ----------
    @staticmethod
    def _extract_solution_from_statevector(
        sv_like, n_phase: int, n_work: int, anc_index: int, total_n: int
    ) -> np.ndarray:
        sv = np.asarray(sv_like, dtype=complex).reshape(-1)
        x_vec = np.zeros(1 << n_work, dtype=complex)
        for idx in range(1 << total_n):
            anc_bit = (idx >> anc_index) & 1
            if anc_bit != 1:
                continue
            w_small = 0
            for j in range(n_work):
                bit = (idx >> j) & 1
                w_small |= (bit << j)
            x_vec[w_small] += sv[idx]
        nrm = np.linalg.norm(x_vec)
        if nrm > 0:
            x_vec = x_vec / nrm
        return x_vec


    # ---------- 运行 ----------
    def run(self, statevector: bool = True) -> HHLResult:
        if self._qc is None:
            self.construct_circuit(measurement=False)

        qc = self._qc
        reg = self._reg or {}
        phase = reg.get("phase", [])
        work  = reg.get("work", [])
        anc   = reg.get("anc", [])[0]
        aux   = reg.get("aux", [])

        # 仅 statevector 路径（cqlib 常用）
        sim = StatevectorSimulator(qc)
        sv_raw = sim.statevector()

        # 总比特 = phase + work + （lookup 的 work/msq/anc/aux）
        # 但我们只需要“实际电路”总比特数。用 qc.num_qubits 最稳妥：
        total_n = qc.num_qubits

        # 统一为 ndarray
        sv_arr = _coerce_statevector_to_array(sv_like=sv_raw, total_n=total_n)

        # anc=1 成功概率
        prob_anc1 = 0.0
        for idx in range(1 << total_n):
            if ((idx >> anc) & 1) == 1:
                prob_anc1 += float(abs(sv_arr[idx]) ** 2)

        # 提取量子方向解（注意 anc_index 用 lookup 里给出的 anc）
        x_quantum = self._extract_solution_from_statevector(
            sv_like=sv_arr,
            n_phase=len(phase),
            n_work=len(work),
            anc_index=anc,
            total_n=total_n
        )

        # 相位对齐后仅比较方向
        x_classic_dir = np.linalg.solve(self.A, self.b)
        x_classic_dir = x_classic_dir / np.linalg.norm(x_classic_dir)
        phase_corr = np.vdot(x_classic_dir, x_quantum)
        if abs(phase_corr) > 0:
            x_quantum *= phase_corr / abs(phase_corr)

        output = x_quantum.copy()

        # 从量子方向解恢复“经典尺度”
        x_aligned = output * np.exp(-1j * np.angle(output[0])) if output.size else output
        x_dir = np.real(x_aligned)
        if np.linalg.norm(x_dir) == 0 and output.size:
            x_dir = np.real(output)
        if np.linalg.norm(x_dir) == 0:
            x_dir = output.real
        if np.linalg.norm(x_dir) > 0:
            x_dir = x_dir / np.linalg.norm(x_dir)

        # 在“电路尺寸”下缩放（与 A,b 一致）
        denom = np.linalg.norm(self.A @ x_dir) if x_dir.size else 0.0
        alpha = (np.linalg.norm(self.b) / denom) if denom > 0 else 0.0
        x_solution_full = alpha * x_dir
        output_full = output

        # 若启用了截断，则回到原尺寸
        x_solution = x_solution_full
        output = output_full
        if self._truncate_hermitian:
            x_solution = x_solution[: self._orig_size]
            output = output[: self._orig_size]
        if self._truncate_powerdim:
            x_solution = x_solution[: self._orig_size]
            output = output[: self._orig_size]

        circuit_info = {
            "num_qubits": qc.num_qubits,
            "num_instructions": getattr(qc, "num_instructions", None)
        }

        return HHLResult(
            solution=x_solution,
            output=output,
            probability_result=float(prob_anc1),
            matrix=self._orig_A,
            vector=self._orig_b,
            circuit_info=circuit_info
        )

# ------------- 简易自测 -------------
if __name__ == "__main__":
    A = np.array([[1, 1/3], [1/3, 1]], dtype=float)
    b = np.array([1.0, 0], dtype=float)

    hhl = HHL(A, b, m_phase=8, trotter_slices=2, truncate_powerdim=True, truncate_hermitian=True)
    qc = hhl.construct_circuit()
    res = hhl.run()

    x_true = np.linalg.solve(A, b)
    rel_err = np.linalg.norm(res.solution - x_true) / np.linalg.norm(x_true)

    print("=== HHL ===")
    print("A =\n", A)
    print("b =", b)
    print("quantum output (direction) =", res.output)
    print("recovered classical solution =", res.solution)
    print("true solution =", x_true)
    print("relative error =", f"{rel_err:.6%}")
    print("success prob anc=1 =", res.probability_result)
    print("circuit info:", res.circuit_info)